import gradio as gr
import numpy as np
from PIL import Image
from pinecone import Pinecone, ServerlessSpec
from collections import Counter
import logging
import time

# -------------------------- 1. 配置带日期的Logging --------------------------
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)
logger = logging.getLogger(__name__)

# -------------------------- 2. Pinecone初始化（含预热优化首次推理） --------------------------
# 替换为你的Pinecone API密钥
PINECONE_API_KEY = "pcsk_2E35Na_3qapzC6AwtSBcZRTjKEczM4NVv8BS69ZSMku7GFtahdRFdovykfLvAzpnXpqojy"
INDEX_NAME = "mnist-index"  # 与之前创建的索引名一致
pc = Pinecone(api_key=PINECONE_API_KEY)
pinecone_index = None

# 连接索引并预热
try:
    # 检查索引是否存在
    if INDEX_NAME in pc.list_indexes().names():
        pinecone_index = pc.Index(INDEX_NAME)
        logger.info(f"✅ 成功连接到Pinecone索引 '{INDEX_NAME}'")
        
        # 预热索引（解决首次推理慢的问题）
        dummy_vector = [0.0] * 64  # 生成64维空向量
        pinecone_index.query(vector=dummy_vector, top_k=1)
        logger.info("✅ Pinecone索引预热完成，首次推理延迟已优化")
        
        # 额外检查索引数据量
        index_stats = pinecone_index.describe_index_stats()
        logger.info(f"🔍 索引当前向量数量: {index_stats.get('vector_count', 0)}")
        if index_stats.get('vector_count', 0) < 100:
            logger.warning("⚠️ 索引向量数量异常（建议至少上千条），请检查数据上传脚本")
    else:
        logger.error(f"❌ 索引 '{INDEX_NAME}' 不存在，请先运行索引创建脚本")
except Exception as e:
    logger.error(f"❌ Pinecone初始化失败: {str(e)}")


# -------------------------- 3. 图像预处理函数（适配64维向量） --------------------------
def preprocess_sketch_image(sketch_data):
    """将手写板图像转为Pinecone查询所需的64维向量"""
    if sketch_data is None:
        raise ValueError("请先在手写板上绘制数字")
    
    # 处理Gradio手写板数据格式
    if isinstance(sketch_data, dict):
        image_array = sketch_data["composite"]
    else:
        image_array = sketch_data
    
    # 转换为PIL图像并处理
    pil_image = Image.fromarray(image_array)
    # 转为灰度图
    if pil_image.mode != 'L':
        pil_image = pil_image.convert('L')
    # 缩放到8×8像素
    resized_image = pil_image.resize((8, 8), Image.Resampling.LANCZOS)
    img_array = np.array(resized_image)
    
    # 颜色反转（手写板白底黑字 → 向量库黑底白字）
    if np.mean(img_array) > 128:
        img_array = 255 - img_array
        logger.info("🔄 执行颜色反转适配向量库格式")
    
    # 归一化到0-16范围（匹配load_digits数据集）
    normalized_array = (img_array / 255.0) * 16
    # 展平为1维列表（Pinecone查询要求）
    flattened_array = normalized_array.flatten().tolist()
    
    # 调试：打印预处理后向量信息
    logger.info(f"🔍 预处理后向量维度: {len(flattened_array)}，前5个值: {flattened_array[:5]}")
    return flattened_array


# -------------------------- 4. Pinecone推理函数（增强错误处理） --------------------------
def predict_handwritten_digit(sketch_input):
    """通过Pinecone向量查询实现手写数字识别，增强错误处理"""
    if pinecone_index is None:
        return "❌ 错误: Pinecone索引未就绪，请检查配置"
    
    try:
        # 1. 预处理图像得到64维查询向量
        query_vector = preprocess_sketch_image(sketch_input)
        
        # 2. Pinecone查询Top-11相似向量（k=11）
        results = pinecone_index.query(
            vector=query_vector,
            top_k=11,
            include_metadata=True  # 包含元数据中的标签
        )
        
        # 调试：打印查询结果
        logger.info(f"🔍 Pinecone查询结果: {results}")
        
        # 3. 校验查询结果是否有效
        if len(results.get('matches', [])) == 0:
            error_msg = "❌ 查询无结果，请检查绘制的数字是否清晰，或索引数据是否正常"
            logger.error(error_msg)
            return error_msg
        
        # 4. 提取相似向量的标签并投票
        top_k_labels = [match['metadata']['label'] for match in results['matches']]
        vote_results = Counter(top_k_labels)
        final_prediction = vote_results.most_common(1)[0][0]
        vote_count = vote_results[final_prediction]
        
        # 5. 构建结果文本
        result_text = f"🎯 识别结果: {final_prediction}\n"
        result_text += f"🗳️ 投票支持: {vote_count}/{len(results['matches'])} 票\n"
        result_text += "\n📊 相似向量详情（Top-11）:\n"
        for i, (match, label) in enumerate(zip(results['matches'], top_k_labels), 1):
            result_text += f"  {i}. 数字 {label}，距离 {match['score']:.4f}\n"
        
        logger.info(f"✅ 预测完成: 数字 {final_prediction}")
        return result_text
        
    except IndexError as e:
        error_msg = f"❌ 索引越界错误（通常因查询结果无标签）: {str(e)}"
        logger.error(error_msg)
        return error_msg
    except Exception as e:
        error_msg = f"❌ 识别错误: {str(e)}"
        logger.error(f"识别错误详情: {e}")
        return error_msg


# -------------------------- 5. 清除手写板函数 --------------------------
def clear_sketchpad():
    """清除手写板内容"""
    logger.info("🧹 手写板已清除")
    return None


# -------------------------- 6. Gradio界面构建（保留原交互逻辑） --------------------------
with gr.Blocks(
    title="手写数字识别系统（Pinecone云版）",
    theme=gr.themes.Soft()
) as demo:
    
    # 标题区域（更新技术说明）
    gr.Markdown("""
    # ✏️ 手写数字识别系统
    ### 基于Pinecone云向量数据库实时识别手写数字 0-9
    """)
    
    with gr.Row(equal_height=False):
        # 左侧：手写绘制区域
        with gr.Column(scale=1):
            gr.Markdown("### 🎨 绘制区域")
            
            sketchpad = gr.Sketchpad(
                label="请在下方手写板绘制数字",
                type="numpy",
                height=280,
                width=280,
                image_mode="L",
                show_download_button=False,
                elem_id="sketchpad"
            )
            
            with gr.Row():
                recognize_btn = gr.Button(
                    "🚀 开始识别",
                    variant="primary",
                    size="lg"
                )
                clear_btn = gr.Button(
                    "🧹 清除重写",
                    variant="secondary",
                    size="lg"
                )
        
        # 右侧：结果显示区域
        with gr.Column(scale=1):
            gr.Markdown("### 📊 识别结果")
            
            output_text = gr.Textbox(
                label="预测结果",
                lines=12,
                placeholder="绘制数字后点击'开始识别'按钮...",
                show_copy_button=True
            )
    
    # 使用说明区域（更新技术支持）
    with gr.Accordion("📖 使用指南", open=False):
        gr.Markdown("""
        ### 使用方法：
        1. **绘制数字**: 在左侧手写板上用鼠标绘制数字 0-9
        2. **开始识别**: 点击"🚀 开始识别"按钮进行识别
        3. **查看结果**: 在右侧查看识别结果和相似向量详情
        4. **重新绘制**: 点击"🧹 清除重写"按钮开始新的识别
        
        ### 绘制技巧：
        - 使用清晰的笔画，避免断断续续
        - 数字尽量写大一些，占据手写板的大部分区域
        - 确保数字完整可见，不要超出边界
        - 使用黑色笔迹在白色背景上绘制
        
        ### 技术支持：
        - 算法：Pinecone向量相似度查询 + 投票机制
        - 索引：mnist-index（预存8×8手写数字数据）
        - 输入：8x8像素灰度图像（自动转换为64维向量）
        """)
    
    # 事件处理
    recognize_btn.click(
        fn=predict_handwritten_digit,
        inputs=sketchpad,
        outputs=output_text
    )
    
    clear_btn.click(
        fn=clear_sketchpad,
        inputs=[],
        outputs=sketchpad
    )


# -------------------------- 7. 主程序 --------------------------
if __name__ == "__main__":
    print("=" * 60)
    print("🎯 手写数字识别系统 - Pinecone云版")
    print("📚 使用技术: Pinecone向量数据库 + 向量相似度查询")
    print("=" * 60)
    
    if pinecone_index is not None:
        print("✅ 系统准备就绪")
        print("🌐 启动Web应用中...")
        print("📍 访问地址: http://localhost:8800")
        print("⏹️  停止服务: 按 Ctrl + C")
        print("-" * 60)
        
        try:
            # 启动应用
            demo.launch(
                server_name="0.0.0.0",
                server_port=8800,
                share=False,
                show_error=True,
                inbrowser=True
            )
        except Exception as e:
            print(f"❌ 应用启动失败: {e}")
            print("💡 解决方案:")
            print("  1. 检查端口8800是否被占用")
            print("  2. 尝试更换端口号（如8801）")
            print("  3. 检查网络连接是否正常")
    else:
        print("❌ 系统启动失败")
        print("\n🔧 请检查以下问题:")
        print(f"  1. Pinecone索引 '{INDEX_NAME}' 是否已创建")
        print("  2. 索引是否已上传8×8手写数字数据")
        print("  3. API密钥是否正确配置")